{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "010a246d-56ec-44b6-98a1-cca3723d784c",
   "metadata": {},
   "source": [
    "# Structured Hierarchical Retrieval\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/query_engine/multi_doc_auto_retrieval/multi_doc_auto_retrieval.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "Doing RAG well over multiple documents is hard. A general framework is given a user query, first select the relevant documents before selecting the content inside.\n",
    "\n",
    "But selecting the documents can be tough - how can we dynamically select documents based on different properties depending on the user query? \n",
    "\n",
    "In this notebook we show you our multi-document RAG architecture:\n",
    "\n",
    "- Represent each document as a concise **metadata** dictionary containing different properties: an extracted summary along with structured metadata.\n",
    "- Store this metadata dictionary as filters within a vector database.\n",
    "- Given a user query, first do **auto-retrieval** - infer the relevant semantic query and the set of filters to query this data (effectively combining text-to-SQL and semantic search)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1d4bb83",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-readers-github\n",
    "%pip install llama-index-vector-stores-weaviate\n",
    "%pip install llama-index-llms-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a599ce1-48b1-44f6-846e-b9d3463635fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index llama-hub"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da8ab861-ea33-4057-8634-b3529c577d29",
   "metadata": {},
   "source": [
    "## Setup and Download Data\n",
    "\n",
    "In this section, we'll load in LlamaIndex Github issues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b49dc9d9-668a-4046-b86a-0ab39d74c9a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12776eb2-c61d-4661-89a5-b7d3be2d8c93",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"GITHUB_TOKEN\"] = \"ghp_...\"\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c423120f-0dfe-4e77-aee5-3325c5d1442e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 100 issues in the repo page 1\n",
      "Resulted in 100 documents\n",
      "Found 100 issues in the repo page 2\n",
      "Resulted in 200 documents\n",
      "Found 100 issues in the repo page 3\n",
      "Resulted in 300 documents\n",
      "Found 64 issues in the repo page 4\n",
      "Resulted in 364 documents\n",
      "No more issues found, stopping\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "from llama_index.readers.github import (\n",
    "    GitHubRepositoryIssuesReader,\n",
    "    GitHubIssuesClient,\n",
    ")\n",
    "\n",
    "github_client = GitHubIssuesClient()\n",
    "loader = GitHubRepositoryIssuesReader(\n",
    "    github_client,\n",
    "    owner=\"run-llama\",\n",
    "    repo=\"llama_index\",\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "orig_docs = loader.load_data()\n",
    "\n",
    "limit = 100\n",
    "\n",
    "docs = []\n",
    "for idx, doc in enumerate(orig_docs):\n",
    "    doc.metadata[\"index_id\"] = int(doc.id_)\n",
    "    if idx >= limit:\n",
    "        break\n",
    "    docs.append(doc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f142129c",
   "metadata": {},
   "source": [
    "## Setup the Vector Store and Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f17643d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import weaviate\n",
    "\n",
    "# cloud\n",
    "auth_config = weaviate.AuthApiKey(\n",
    "    api_key=\"XRa15cDIkYRT7AkrpqT6jLfE4wropK1c1TGk\"\n",
    ")\n",
    "client = weaviate.Client(\n",
    "    \"https://llama-index-test-v0oggsoz.weaviate.network\",\n",
    "    auth_client_secret=auth_config,\n",
    ")\n",
    "\n",
    "class_name = \"LlamaIndex_docs\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75615d58",
   "metadata": {},
   "outputs": [],
   "source": [
    "# optional: delete schema\n",
    "client.schema.delete_class(class_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b67acd54",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.vector_stores.weaviate import WeaviateVectorStore\n",
    "from llama_index.core import VectorStoreIndex, StorageContext\n",
    "\n",
    "vector_store = WeaviateVectorStore(\n",
    "    weaviate_client=client, index_name=class_name\n",
    ")\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82d8ffb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_index = VectorStoreIndex.from_documents(\n",
    "    docs, storage_context=storage_context\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cb80399",
   "metadata": {},
   "source": [
    "## Create IndexNodes for retrieval and filtering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee7f89b0-4c55-4bfe-83c8-8539b7939de2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SummaryIndex\n",
    "from llama_index.core.async_utils import run_jobs\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core.schema import IndexNode\n",
    "from llama_index.core.vector_stores import (\n",
    "    FilterOperator,\n",
    "    MetadataFilter,\n",
    "    MetadataFilters,\n",
    ")\n",
    "\n",
    "\n",
    "async def aprocess_doc(doc, include_summary: bool = True):\n",
    "    \"\"\"Process doc.\"\"\"\n",
    "    metadata = doc.metadata\n",
    "\n",
    "    date_tokens = metadata[\"created_at\"].split(\"T\")[0].split(\"-\")\n",
    "    year = int(date_tokens[0])\n",
    "    month = int(date_tokens[1])\n",
    "    day = int(date_tokens[2])\n",
    "\n",
    "    assignee = (\n",
    "        \"\" if \"assignee\" not in doc.metadata else doc.metadata[\"assignee\"]\n",
    "    )\n",
    "    size = \"\"\n",
    "    if len(doc.metadata[\"labels\"]) > 0:\n",
    "        size_arr = [l for l in doc.metadata[\"labels\"] if \"size:\" in l]\n",
    "        size = size_arr[0].split(\":\")[1] if len(size_arr) > 0 else \"\"\n",
    "    new_metadata = {\n",
    "        \"state\": metadata[\"state\"],\n",
    "        \"year\": year,\n",
    "        \"month\": month,\n",
    "        \"day\": day,\n",
    "        \"assignee\": assignee,\n",
    "        \"size\": size,\n",
    "    }\n",
    "\n",
    "    # now extract out summary\n",
    "    summary_index = SummaryIndex.from_documents([doc])\n",
    "    query_str = \"Give a one-sentence concise summary of this issue.\"\n",
    "    query_engine = summary_index.as_query_engine(\n",
    "        llm=OpenAI(model=\"gpt-3.5-turbo\")\n",
    "    )\n",
    "    summary_txt = await query_engine.aquery(query_str)\n",
    "    summary_txt = str(summary_txt)\n",
    "\n",
    "    index_id = doc.metadata[\"index_id\"]\n",
    "    # filter for the specific doc id\n",
    "    filters = MetadataFilters(\n",
    "        filters=[\n",
    "            MetadataFilter(\n",
    "                key=\"index_id\", operator=FilterOperator.EQ, value=int(index_id)\n",
    "            ),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    # create an index node using the summary text\n",
    "    index_node = IndexNode(\n",
    "        text=summary_txt,\n",
    "        metadata=new_metadata,\n",
    "        obj=doc_index.as_retriever(filters=filters),\n",
    "        index_id=doc.id_,\n",
    "    )\n",
    "\n",
    "    return index_node\n",
    "\n",
    "\n",
    "async def aprocess_docs(docs):\n",
    "    \"\"\"Process metadata on docs.\"\"\"\n",
    "\n",
    "    index_nodes = []\n",
    "    tasks = []\n",
    "    for doc in docs:\n",
    "        task = aprocess_doc(doc)\n",
    "        tasks.append(task)\n",
    "\n",
    "    index_nodes = await run_jobs(tasks, show_progress=True, workers=3)\n",
    "\n",
    "    return index_nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17728251-e7c8-47eb-b139-ee0a7246f894",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|          | 1/100 [00:00<00:55,  1.78it/s]/home/loganm/llama_index_proper/llama_index/.venv/lib/python3.11/site-packages/openai/_resource.py:38: ResourceWarning: unclosed <socket.socket fd=71, family=2, type=1, proto=6, laddr=('172.25.21.0', 40832), raddr=('104.18.7.192', 443)>\n",
      "  self._delete = client.delete\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=73 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=71 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      " 12%|█▏        | 12/100 [00:04<00:31,  2.79it/s]/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=76 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=77 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=78 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/llama_index_proper/llama_index/.venv/lib/python3.11/site-packages/openai/resources/chat/completions.py:1337: ResourceWarning: unclosed <socket.socket fd=81, family=2, type=1, proto=6, laddr=('172.25.21.0', 40848), raddr=('104.18.7.192', 443)>\n",
      "  completions.create,\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=81 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=82 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=83 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=84 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      " 21%|██        | 21/100 [00:06<00:22,  3.58it/s]/home/loganm/llama_index_proper/llama_index/.venv/lib/python3.11/site-packages/openai/_resource.py:34: ResourceWarning: unclosed <socket.socket fd=81, family=2, type=1, proto=6, laddr=('172.25.21.0', 40866), raddr=('104.18.7.192', 443)>\n",
      "  self._get = client.get\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/llama_index_proper/llama_index/.venv/lib/python3.11/site-packages/openai/_resource.py:34: ResourceWarning: unclosed <socket.socket fd=82, family=2, type=1, proto=6, laddr=('172.25.21.0', 40868), raddr=('104.18.7.192', 443)>\n",
      "  self._get = client.get\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=86 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      " 38%|███▊      | 38/100 [00:12<00:24,  2.54it/s]/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=90 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=92 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/llama_index_proper/llama_index/.venv/lib/python3.11/site-packages/openai/_resource.py:34: ResourceWarning: unclosed <socket.socket fd=94, family=2, type=1, proto=6, laddr=('172.25.21.0', 40912), raddr=('104.18.7.192', 443)>\n",
      "  self._get = client.get\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=94 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      " 50%|█████     | 50/100 [00:17<00:19,  2.51it/s]/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=95 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=96 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=97 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      " 73%|███████▎  | 73/100 [00:24<00:07,  3.42it/s]/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=101 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=102 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      " 82%|████████▏ | 82/100 [00:27<00:06,  2.94it/s]/home/loganm/miniconda3/envs/llama_index/lib/python3.11/functools.py:76: ResourceWarning: unclosed <socket.socket fd=102, family=2, type=1, proto=6, laddr=('172.25.21.0', 40998), raddr=('104.18.7.192', 443)>\n",
      "  return partial(update_wrapper, wrapped=wrapped,\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      " 92%|█████████▏| 92/100 [00:32<00:03,  2.15it/s]/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=106 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "/home/loganm/miniconda3/envs/llama_index/lib/python3.11/asyncio/selector_events.py:835: ResourceWarning: unclosed transport <_SelectorSocketTransport fd=111 read=idle write=<idle, bufsize=0>>\n",
      "  _warn(f\"unclosed transport {self!r}\", ResourceWarning, source=self)\n",
      "ResourceWarning: Enable tracemalloc to get the object allocation traceback\n",
      "100%|██████████| 100/100 [00:36<00:00,  2.71it/s]\n"
     ]
    }
   ],
   "source": [
    "index_nodes = await aprocess_docs(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6907607-47b7-4966-9501-6c5320ec66e5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'state': 'open',\n",
       " 'year': 2024,\n",
       " 'month': 1,\n",
       " 'day': 13,\n",
       " 'assignee': '',\n",
       " 'size': 'XL'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index_nodes[5].metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "773585a7-3027-4a12-9349-7320822514e0",
   "metadata": {},
   "source": [
    "## Create the Top-Level AutoRetriever\n",
    "\n",
    "We load both the summarized metadata as well as the original docs into the vector database.\n",
    "1. **Summarized Metadata**: This goes into the `LlamaIndex_auto` collection.\n",
    "2. **Original Docs**: This goes into the `LlamaIndex_docs` collection.\n",
    "\n",
    "By storing both the summarized metadata as well as the original documents, we can execute our structured, hierarchical retrieval strategies.\n",
    "\n",
    "We load into a vector database that supports auto-retrieval. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55c7baf6-4267-4543-8a2c-7b44a8fcd017",
   "metadata": {},
   "source": [
    "### Load Summarized Metadata\n",
    "\n",
    "This goes into `LlamaIndex_auto`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fcbe94d-0fe4-48b1-954c-31d0f278ebc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import weaviate\n",
    "\n",
    "# cloud\n",
    "auth_config = weaviate.AuthApiKey(\n",
    "    api_key=\"XRa15cDIkYRT7AkrpqT6jLfE4wropK1c1TGk\"\n",
    ")\n",
    "client = weaviate.Client(\n",
    "    \"https://llama-index-test-v0oggsoz.weaviate.network\",\n",
    "    auth_client_secret=auth_config,\n",
    ")\n",
    "\n",
    "class_name = \"LlamaIndex_auto\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "782b544b-2e8e-4fde-a35f-fb77133a0cef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# optional: delete schema\n",
    "client.schema.delete_class(class_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd78a330",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.vector_stores.weaviate import WeaviateVectorStore\n",
    "from llama_index.core import VectorStoreIndex, StorageContext\n",
    "\n",
    "vector_store_auto = WeaviateVectorStore(\n",
    "    weaviate_client=client, index_name=class_name\n",
    ")\n",
    "storage_context_auto = StorageContext.from_defaults(\n",
    "    vector_store=vector_store_auto\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87910444-2cfb-4c47-8821-83d679486700",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Since \"index_nodes\" are concise summaries, we can directly feed them as objects into VectorStoreIndex\n",
    "index = VectorStoreIndex(\n",
    "    objects=index_nodes, storage_context=storage_context_auto\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f44433ac-3a4c-4bda-bb5b-6bbfd18c1ddf",
   "metadata": {},
   "source": [
    "## Setup Composable Auto-Retriever\n",
    "\n",
    "In this section we setup our auto-retriever. There's a few steps that we need to perform.\n",
    "\n",
    "1. **Define the Schema**: Define the vector db schema (e.g. the metadata fields). This will be put into the LLM input prompt when it's deciding what metadata filters to infer.\n",
    "2. **Instantiate the VectorIndexAutoRetriever class**: This creates a retriever on top of our summarized metadata index, and takes in the defined schema as input.\n",
    "3. **Define a wrapper retriever**: This allows us to postprocess each node into an `IndexNode`, with an index id linking back source document. This will allow us to do recursive retrieval in the next section (which depends on IndexNode objects linking to downstream retrievers/query engines/other Nodes). **NOTE**: We are working on improving this abstraction.\n",
    "\n",
    "Running this retriever will retrieve based on our text summaries and metadat of our top-level `IndeNode` objects. Then, their underlying retrievers will be used to retrieve content from the specific github issue."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c926128-c1aa-4f42-8061-b3f2df1272d0",
   "metadata": {},
   "source": [
    "### 1. Define the Schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d01aade-d676-49db-b851-2ac62b4e53c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.vector_stores import MetadataInfo, VectorStoreInfo\n",
    "\n",
    "\n",
    "vector_store_info = VectorStoreInfo(\n",
    "    content_info=\"Github Issues\",\n",
    "    metadata_info=[\n",
    "        MetadataInfo(\n",
    "            name=\"state\",\n",
    "            description=\"Whether the issue is `open` or `closed`\",\n",
    "            type=\"string\",\n",
    "        ),\n",
    "        MetadataInfo(\n",
    "            name=\"year\",\n",
    "            description=\"The year issue was created\",\n",
    "            type=\"integer\",\n",
    "        ),\n",
    "        MetadataInfo(\n",
    "            name=\"month\",\n",
    "            description=\"The month issue was created\",\n",
    "            type=\"integer\",\n",
    "        ),\n",
    "        MetadataInfo(\n",
    "            name=\"day\",\n",
    "            description=\"The day issue was created\",\n",
    "            type=\"integer\",\n",
    "        ),\n",
    "        MetadataInfo(\n",
    "            name=\"assignee\",\n",
    "            description=\"The assignee of the ticket\",\n",
    "            type=\"string\",\n",
    "        ),\n",
    "        MetadataInfo(\n",
    "            name=\"size\",\n",
    "            description=\"How big the issue is (XS, S, M, L, XL, XXL)\",\n",
    "            type=\"string\",\n",
    "        ),\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6a69271-e597-40c6-88aa-a755bfd75754",
   "metadata": {},
   "source": [
    "### 2. Instantiate VectorIndexAutoRetriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9125a832-940a-44f3-be91-ad17cdfc267b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.retrievers import VectorIndexAutoRetriever\n",
    "\n",
    "retriever = VectorIndexAutoRetriever(\n",
    "    index,\n",
    "    vector_store_info=vector_store_info,\n",
    "    similarity_top_k=2,\n",
    "    empty_query_top_k=10,  # if only metadata filters are specified, this is the limit\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97a91a51-9a2c-447e-9848-53ee0a705baa",
   "metadata": {},
   "source": [
    "## Try It Out\n",
    "\n",
    "Now we can start retrieving relevant context over Github Issues! \n",
    "\n",
    "To complete the RAG pipeline setup we'll combine our recursive retriever with our `RetrieverQueryEngine` to generate a response in addition to the retrieved nodes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a2032ef-e776-49d8-b8d4-f8b79fbd3599",
   "metadata": {},
   "source": [
    "### Try Out Retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7136a40b-c615-4f32-bc98-fdc3f72f3085",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using query str: issues\n",
      "Using filters: [('day', '==', '11'), ('month', '==', '01')]\n",
      "\u001b[1;3;38;2;11;159;203mRetrieval entering 9995: VectorIndexRetriever\n",
      "\u001b[0m\u001b[1;3;38;2;237;90;200mRetrieving from object VectorIndexRetriever with query issues\n",
      "\u001b[0m\u001b[1;3;38;2;11;159;203mRetrieval entering 9985: VectorIndexRetriever\n",
      "\u001b[0m\u001b[1;3;38;2;237;90;200mRetrieving from object VectorIndexRetriever with query issues\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "from llama_index.core import QueryBundle\n",
    "\n",
    "nodes = retriever.retrieve(QueryBundle(\"Tell me about some issues on 01/11\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "002f3135-6b88-45fa-a25d-1b4e7e977489",
   "metadata": {},
   "source": [
    "The result is the source chunks in the relevant docs. \n",
    "\n",
    "Let's look at the date attached to the source chunk (was present in the original metadata)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3804c74f-2bb5-4935-b15b-ff16ce0a7475",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of source nodes: 2\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'state': 'open',\n",
       " 'created_at': '2024-01-11T20:37:34Z',\n",
       " 'url': 'https://api.github.com/repos/run-llama/llama_index/issues/9995',\n",
       " 'source': 'https://github.com/run-llama/llama_index/pull/9995',\n",
       " 'labels': ['size:XXL'],\n",
       " 'index_id': 9995}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(f\"Number of source nodes: {len(nodes)}\")\n",
    "nodes[0].node.metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97cba7dd-b5f5-4759-9dcf-b9ee06a7ec29",
   "metadata": {},
   "source": [
    "### Plug into `RetrieverQueryEngine`\n",
    "\n",
    "We plug into RetrieverQueryEngine to synthesize a result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9bdaf13-71b4-43c4-b49c-8c9a109819f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.query_engine import RetrieverQueryEngine\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "\n",
    "query_engine = RetrieverQueryEngine.from_args(retriever, llm=llm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10d52659-c297-458c-aba6-496995655ba2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using query str: issues\n",
      "Using filters: [('day', '==', '11'), ('month', '==', '01')]\n",
      "\u001b[1;3;38;2;11;159;203mRetrieval entering 9995: VectorIndexRetriever\n",
      "\u001b[0m\u001b[1;3;38;2;237;90;200mRetrieving from object VectorIndexRetriever with query issues\n",
      "\u001b[0m\u001b[1;3;38;2;11;159;203mRetrieval entering 9985: VectorIndexRetriever\n",
      "\u001b[0m\u001b[1;3;38;2;237;90;200mRetrieving from object VectorIndexRetriever with query issues\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\"Tell me about some issues on 01/11\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f591dae0-0438-4018-9732-c5aa9357938a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are two issues that were created on 01/11. The first issue is related to ensuring backwards compatibility with the new Pinecone client version bifurcation. The second issue is a feature request to implement the Language Agent Tree Search (LATS) agent in llama-index.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6450df09-f8c4-4dee-aa10-6c85c0ea362b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using query str: agents\n",
      "Using filters: [('state', '==', 'open')]\n",
      "\u001b[1;3;38;2;11;159;203mRetrieval entering 10058: VectorIndexRetriever\n",
      "\u001b[0m\u001b[1;3;38;2;237;90;200mRetrieving from object VectorIndexRetriever with query agents\n",
      "\u001b[0m\u001b[1;3;38;2;11;159;203mRetrieval entering 9899: VectorIndexRetriever\n",
      "\u001b[0m\u001b[1;3;38;2;237;90;200mRetrieving from object VectorIndexRetriever with query agents\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\n",
    "    \"Tell me about some open issues related to agents\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb0d36e2-f4c4-4cca-a4b9-ddd4f6d2e9a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are two open issues related to agents. One issue is about adding context for agents, updating a stale link, and adding a notebook to demo a react agent with context. The other issue is a feature request for parallelism when using the top agent from a multi-document agent while comparing multiple documents.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "230a7945-0a77-4d0c-a324-f24e2398c146",
   "metadata": {},
   "source": [
    "## Concluding Thoughts\n",
    "\n",
    "This shows you how to create a structured retrieval layer over your document summaries, allowing you to dynamically pull in the relevant documents based on the user query.\n",
    "\n",
    "You may notice similarities between this and our [multi-document agents](https://docs.llamaindex.ai/en/stable/examples/agent/multi_document_agents.html). Both architectures are aimed for powerful multi-document retrieval.\n",
    "\n",
    "The goal of this notebook is to show you how to apply structured querying in a multi-document setting. You can actually apply this auto-retrieval algorithm to our multi-agent setup too. The multi-agent setup is primarily focused on adding agentic reasoning across documents and per documents, alloinwg multi-part queries using chain-of-thought."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
